{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 设计模式\n",
    "\n",
    "### 基于单例工厂、策略模式的多知识库服务"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# -*- coding: utf-8 -/-\n",
    "# @Author       :  lizhenping\n",
    "# @Project_Name : XXX\n",
    "# @FileName     : XXX.py\n",
    "# @SoftWare     : Vsocode\n",
    "# @Time         : 2024/6/12 12:46\n",
    "\n",
    "\n",
    "from abc import ABC, abstractmethod\n",
    "\n",
    "class KBServiceFactory:\n",
    "    @staticmethod\n",
    "    def create_service(service_type):\n",
    "        \"\"\"\n",
    "        创建知识库服务的工厂方法\n",
    "        :param service_type: 服务类型\n",
    "        :return: 知识库服务实例\n",
    "        \"\"\"\n",
    "        if service_type == \"policy\":\n",
    "            return PolicyKBService()\n",
    "        elif service_type == \"report\":\n",
    "            return ReportKBService()\n",
    "        elif service_type == \"journal\":\n",
    "            return JournalKBService()\n",
    "        else:\n",
    "            raise ValueError(f\"Unknown service type: {service_type}\")\n",
    "\n",
    "# query_rewrite_strategy.py\n",
    "class QueryRewriteStrategy(ABC):\n",
    "    @abstractmethod\n",
    "    def rewrite(self, query, history):\n",
    "        \"\"\"\n",
    "        查询改写策略接口\n",
    "        :param query: 原始查询\n",
    "        :param history: 历史记录\n",
    "        :return: 改写后的查询\n",
    "        \"\"\"\n",
    "        pass\n",
    "\n",
    "class PolicyQueryRewriteStrategy(QueryRewriteStrategy):\n",
    "    def rewrite(self, query, history):\n",
    "        \"\"\"\n",
    "        XX查询改写策略\n",
    "        \"\"\"\n",
    "        rewritten_query = get_llm_model_response(\n",
    "            strategy_name=\"query rewrite\",\n",
    "            llm_model_name=\"query_rewrite_model\",\n",
    "            template_prompt_name=\"query_rewrite_policy\",\n",
    "            prompt_param_dict={\"query\": query, \"history\": history, \"time\": current_time},\n",
    "            temperature=0.01,\n",
    "            max_tokens=512\n",
    "        )\n",
    "        return rewritten_query\n",
    "\n",
    "class ReportQueryRewriteStrategy(QueryRewriteStrategy):\n",
    "    def rewrite(self, query, history):\n",
    "        \"\"\"\n",
    "        XX查询改写策略\n",
    "        \"\"\"\n",
    "        rewritten_query = get_llm_model_response(\n",
    "            strategy_name=\"query rewrite\",\n",
    "            llm_model_name=\"query_rewrite_model\",\n",
    "            template_prompt_name=\"query_rewrite_report\",\n",
    "            prompt_param_dict={\"query\": query, \"history\": history},\n",
    "            temperature=0.01,\n",
    "            max_tokens=512\n",
    "        )\n",
    "        return rewritten_query\n",
    "\n",
    "class JournalQueryRewriteStrategy(QueryRewriteStrategy):\n",
    "    def rewrite(self, query, history):\n",
    "        \"\"\"\n",
    "        XX查询改写策略\n",
    "        \"\"\"\n",
    "        rewritten_query = get_llm_model_response(\n",
    "            strategy_name=\"query rewrite\",\n",
    "            llm_model_name=\"query_rewrite_model\",\n",
    "            template_prompt_name=\"query_rewrite\",\n",
    "            prompt_param_dict={\"query\": query, \"history\": history},\n",
    "            temperature=0.01,\n",
    "            max_tokens=512\n",
    "        )\n",
    "        return rewritten_query\n",
    "\n",
    "# summary_strategy.py\n",
    "class SummaryStrategy(ABC):\n",
    "    @abstractmethod\n",
    "    def summarize(self, docs, chunk_size, min_chunk_size, summary_model_name):\n",
    "        \"\"\"\n",
    "        摘要生成策略接口\n",
    "        :param docs: 文档列表\n",
    "        :param chunk_size: 分块大小\n",
    "        :param min_chunk_size: 最小分块大小\n",
    "        :param summary_model_name: 摘要生成模型名称\n",
    "        :return: 生成的摘要\n",
    "        \"\"\"\n",
    "        pass\n",
    "\n",
    "class TextRankSummaryStrategy(SummaryStrategy):\n",
    "    def summarize(self, docs, chunk_size, min_chunk_size, summary_model_name):\n",
    "        \"\"\"\n",
    "        TextRank摘要生成策略\n",
    "        \"\"\"\n",
    "        summary = \"\"\n",
    "        for doc in docs:\n",
    "            summary += doc.page_content\n",
    "        return TextRank(summary, 240)\n",
    "\n",
    "# kb_service.py\n",
    "class KBService(ABC):\n",
    "    def __init__(self, query_rewrite_strategy, summary_strategy):\n",
    "        \"\"\"\n",
    "        知识库服务基类\n",
    "        :param query_rewrite_strategy: 查询改写策略\n",
    "        :param summary_strategy: 摘要生成策略\n",
    "        \"\"\"\n",
    "        self.query_rewrite_strategy = query_rewrite_strategy\n",
    "        self.summary_strategy = summary_strategy\n",
    "\n",
    "    @abstractmethod\n",
    "    def search_docs(self, query, top_k, score_threshold):\n",
    "        \"\"\"\n",
    "        搜索文档接口\n",
    "        :param query: 查询\n",
    "        :param top_k: 返回的文档数量\n",
    "        :param score_threshold: 相关性得分阈值\n",
    "        :return: 搜索结果文档列表\n",
    "        \"\"\"\n",
    "        pass\n",
    "\n",
    "    def process_query(self, query, history, top_k, score_threshold, chunk_size, min_chunk_size, summary_model_name):\n",
    "        \"\"\"\n",
    "        处理查询的模板方法\n",
    "        :param query: 查询\n",
    "        :param history: 历史记录\n",
    "        :param top_k: 返回的文档数量\n",
    "        :param score_threshold: 相关性得分阈值\n",
    "        :param chunk_size: 分块大小\n",
    "        :param min_chunk_size: 最小分块大小\n",
    "        :param summary_model_name: 摘要生成模型名称\n",
    "        :return: 摘要和搜索结果文档列表\n",
    "        \"\"\"\n",
    "        rewritten_query = self.query_rewrite_strategy.rewrite(query, history)\n",
    "        docs = self.search_docs(rewritten_query, top_k, score_threshold)\n",
    "        summary = self.summary_strategy.summarize(docs, chunk_size, min_chunk_size, summary_model_name)\n",
    "        return summary, docs\n",
    "\n",
    "class PolicyKBService(KBService):\n",
    "    def __init__(self):\n",
    "        \"\"\"\n",
    "        XX知识库服务\n",
    "        \"\"\"\n",
    "        super().__init__(PolicyQueryRewriteStrategy(), TextRankSummaryStrategy())\n",
    "\n",
    "    def search_docs(self, query, top_k, score_threshold):\n",
    "        \"\"\"\n",
    "        搜索XX文档\n",
    "        \"\"\"\n",
    "        return search_docs_return_score(\n",
    "            fileName=[],\n",
    "            query=query,\n",
    "            usr_query=query,\n",
    "            knowledge_base_name=POLICY_KNOWLEDGE_BASE,\n",
    "            top_k=top_k,\n",
    "            score_threshold=score_threshold\n",
    "        )\n",
    "\n",
    "class ReportKBService(KBService):\n",
    "    def __init__(self):\n",
    "        \"\"\"\n",
    "        XX知识库服务\n",
    "        \"\"\"\n",
    "        super().__init__(ReportQueryRewriteStrategy(), TextRankSummaryStrategy())\n",
    "\n",
    "    def search_docs(self, query, top_k, score_threshold):\n",
    "        \"\"\"\n",
    "        搜索XX文档\n",
    "        \"\"\"\n",
    "        return search_docs(\n",
    "            fileName=[],\n",
    "            query=query,\n",
    "            knowledge_base_name=REPORT_KNOWLEDGE_BASE,\n",
    "            top_k=top_k,\n",
    "            score_threshold=score_threshold,\n",
    "            expr=\" _type == 'content'\"\n",
    "        )\n",
    "\n",
    "class JournalKBService(KBService):\n",
    "    def __init__(self):\n",
    "        \"\"\"\n",
    "        XX知识库服务\n",
    "        \"\"\"\n",
    "        super().__init__(JournalQueryRewriteStrategy(), TextRankSummaryStrategy())\n",
    "\n",
    "    def search_docs(self, query, top_k, score_threshold):\n",
    "        \"\"\"\n",
    "        搜索XX文档\n",
    "        \"\"\"\n",
    "        return search_docs(\n",
    "            fileName=[],\n",
    "            query=query,\n",
    "            knowledge_base_name=JOURNAL_KNOWLEDGE_BASE,\n",
    "            top_k=top_k,\n",
    "            score_threshold=score_threshold\n",
    "        )\n",
    "\n",
    "# prompt_builder.py\n",
    "from langchain.prompts.chat import ChatPromptTemplate\n",
    "from langchain import PromptTemplate\n",
    "\n",
    "class PromptBuilder:\n",
    "    def __init__(self):\n",
    "        \"\"\"\n",
    "        提示构建器\n",
    "        \"\"\"\n",
    "        self.prompt_parts = []\n",
    "\n",
    "    def add_system_message(self, content):\n",
    "        \"\"\"\n",
    "        添加系统消息\n",
    "        :param content: 消息内容\n",
    "        :return: self\n",
    "        \"\"\"\n",
    "        self.prompt_parts.append({\"role\": \"system\", \"content\": content})\n",
    "        return self\n",
    "\n",
    "    def add_user_message(self, content):\n",
    "        \"\"\"\n",
    "        添加用户消息\n",
    "        :param content: 消息内容\n",
    "        :return: self\n",
    "        \"\"\"\n",
    "        self.prompt_parts.append({\"role\": \"user\", \"content\": content})\n",
    "        return self\n",
    "\n",
    "    def add_history(self, history):\n",
    "        \"\"\"\n",
    "        添加历史记录\n",
    "        :param history: 历史记录列表\n",
    "        :return: self\n",
    "        \"\"\"\n",
    "        self.prompt_parts.extend([h.to_msg_template() for h in history])\n",
    "        return self\n",
    "\n",
    "    def build_chat_prompt(self, prompt_template):\n",
    "        \"\"\"\n",
    "        构建聊天提示\n",
    "        :param prompt_template: 提示模板\n",
    "        :return: 聊天提示对象\n",
    "        \"\"\"\n",
    "        input_msg = History(role=\"user\", content=prompt_template).to_msg_template(False)\n",
    "        return ChatPromptTemplate.from_messages(self.prompt_parts + [input_msg])\n",
    "\n",
    "    def build_prompt_template(self, prompt_template):\n",
    "        \"\"\"\n",
    "        构建提示模板\n",
    "        :param prompt_template: 提示模板字符串\n",
    "        :return: 提示模板对象\n",
    "        \"\"\"\n",
    "        return PromptTemplate.from_template(prompt_template)\n",
    "\n",
    "# doc_adapter.py\n",
    "from abc import ABC, abstractmethod\n",
    "\n",
    "class DocAdapter(ABC):\n",
    "    @abstractmethod\n",
    "    def to_source_text(self, doc):\n",
    "        \"\"\"\n",
    "        将文档转换为源文本的接口\n",
    "        :param doc: 文档对象\n",
    "        :return: 源文本\n",
    "        \"\"\"\n",
    "        pass\n",
    "\n",
    "class PolicyDocAdapter(DocAdapter):\n",
    "    def to_source_text(self, doc, inum):\n",
    "        \"\"\"\n",
    "        将XX文档转换为源文本\n",
    "        \"\"\"\n",
    "        filename = doc.metadata.get(\"title\")\n",
    "        detail_url = \"https://policy.xxx.cn/detail/\" + doc.metadata.get(\"primary_key\") + \".html\"\n",
    "        if filename:\n",
    "            if doc.metadata.get('_type') == 'title':\n",
    "                return f\"XX: [{len(source_documents) + 1}][{filename}]({detail_url})\\n\\n\"\n",
    "            else:\n",
    "                return f\"XX: [{len(source_documents) + 1}][{filename}]({detail_url})\\n\\n{doc.page_content}\\n\\n\"\n",
    "        else:\n",
    "            if doc.metadata.get('_type') == 'title':\n",
    "                return f\"XX: [{len(source_documents) + 1}][{\"原文地址\"}]({detail_url})\\n\\n\"\n",
    "            else:\n",
    "                return f\"XX: [{len(source_documents) + 1}][{\"原文地址\"}]({detail_url})\\n\\n{doc.page_content}\\n\\n\"\n",
    "\n",
    "class ReportDocAdapter(DocAdapter):\n",
    "    def to_source_text(self, doc, inum):\n",
    "        \"\"\"\n",
    "        将XX文档转换为源文本\n",
    "        \"\"\"\n",
    "        return f\"XX：[{len(source_documents) + 1}] [{doc.metadata.get('source').replace('.pdf', '')}]()\\n\\n{doc.page_content}\\n\\n\"\n",
    "\n",
    "class JournalDocAdapter(DocAdapter):\n",
    "    def to_source_text(self, doc, inum):\n",
    "        \"\"\"\n",
    "        将XX文档转换为源文本\n",
    "        \"\"\"\n",
    "        return f\"XX论文：[{len(source_documents) + 1}] [{doc.metadata.get('title')}](https://xxx.xxx.cn/xxx/detail/1002/dw_journal_article_20210417/{doc.metadata.get('ID')}.html)\\n\\n\"\n",
    "\n",
    "# main.py\n",
    "async def knowledge_base_chat(query, fileName, knowledge_base_name_list, top_k, score_threshold, history, stream, model_name,\n",
    "                  temperature, max_tokens, prompt_name, request, use_summary, use_model_self_response, chunk_size,\n",
    "                  min_chunk_size, summary_model_name, query_rewrite_model_name):\n",
    "    \"\"\"\n",
    "    知识库聊天主函数\n",
    "    \"\"\"\n",
    "    ...\n",
    "\n",
    "    # 使用工厂创建知识库服务\n",
    "    policy_service = KBServiceFactory.create_service(\"policy\")\n",
    "    report_service = KBServiceFactory.create_service(\"report\")\n",
    "    journal_service = KBServiceFactory.create_service(\"journal\")\n",
    "\n",
    "    # 使用模板方法处理查询\n",
    "    policy_summary, policy_docs = await run_in_threadpool(\n",
    "        policy_service.process_query,\n",
    "        query=query,\n",
    "        history=history,\n",
    "        top_k=top_k,\n",
    "        score_threshold=score_threshold,\n",
    "        chunk_size=chunk_size,\n",
    "        min_chunk_size=min_chunk_size,\n",
    "        summary_model_name=summary_model_name\n",
    "    )\n",
    "    report_summary, report_docs = await run_in_threadpool(\n",
    "        report_service.process_query,\n",
    "        query=query,\n",
    "        history=history,\n",
    "        top_k=top_k,\n",
    "        score_threshold=score_threshold,\n",
    "        chunk_size=chunk_size,\n",
    "        min_chunk_size=min_chunk_size,\n",
    "        summary_model_name=summary_model_name\n",
    "    )\n",
    "    journal_summary, journal_docs = await run_in_threadpool(\n",
    "        journal_service.process_query,\n",
    "        query=query,\n",
    "        history=history,\n",
    "        top_k=top_k,\n",
    "        score_threshold=score_threshold,\n",
    "        chunk_size=chunk_size,\n",
    "        min_chunk_size=min_chunk_size,\n",
    "        summary_model_name=summary_model_name\n",
    "    )\n",
    "\n",
    "    # 使用构建器构建提示\n",
    "    prompt_builder = PromptBuilder()\n",
    "\n",
    "    if use_model_self_response:\n",
    "        # 添加模型自身回答到提示中\n",
    "        model_self_response = get_llm_model_response(\n",
    "            strategy_name=\"self response\",\n",
    "            llm_model_name=query_rewrite_model_name,\n",
    "            template_prompt_name=\"self_response\",\n",
    "            prompt_param_dict={\"query\": query},\n",
    "            temperature=0.01,\n",
    "            max_tokens=512\n",
    "        )\n",
    "        prompt_builder.add_system_message(f\"参考资料[{1}] {model_self_response}\")\n",
    "\n",
    "    prompt_builder.add_system_message(policy_summary)\n",
    "    prompt_builder.add_history(history)\n",
    "    prompt_builder.add_user_message(query)\n",
    "\n",
    "    if history and prompt_name not in [\"Question Assistant\"]:\n",
    "        prompt = prompt_builder.build_chat_prompt(get_prompt_template(\"knowledge_base_chat\", prompt_name))\n",
    "    else:\n",
    "        prompt = prompt_builder.build_prompt_template(get_prompt_template(\"knowledge_base_chat\", prompt_name))\n",
    "\n",
    "    # 使用适配器将文档转换为源文本\n",
    "    policy_doc_adapter = PolicyDocAdapter()\n",
    "    report_doc_adapter = ReportDocAdapter()\n",
    "    journal_doc_adapter = JournalDocAdapter()\n",
    "\n",
    "    source_documents = []\n",
    "\n",
    "    for inum, doc in enumerate(policy_docs):\n",
    "        source_documents.append(policy_doc_adapter.to_source_text(doc, inum))\n",
    "\n",
    "    for inum, doc in enumerate(report_docs):\n",
    "        source_documents.append(report_doc_adapter.to_source_text(doc, inum))\n",
    "\n",
    "    for inum, doc in enumerate(journal_docs):\n",
    "        source_documents.append(journal_doc_adapter.to_source_text(doc, inum))\n",
    "\n",
    "# 遵循高内聚，低耦合。开闭原则。\n",
    "# 其中包含了工厂模式，测策略模式，模板方法，构建起模式，适配器模式。\n",
    "# 代码基于于开源知识库工具 Langchain ChatChat 修改\n",
    "\n",
    "# 设计隐私删除相关数据内容\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
